const std = @import("std");
const Build = std.Build;
const Step = Build.Step;
const Target = std.Target;

const os = @import("kernel/os.zig");

const String = []const u8;
const OsArch = enum { riscv64, riscv32 };

const features = Target.riscv.featureSet(&.{.zihintpause});
const os_arch_specs: std.EnumArray(OsArch, Target.Query) = .init(.{
    .riscv64 = .{ .cpu_arch = .riscv64, .os_tag = .other, .cpu_features_add = features },
    .riscv32 = .{ .cpu_arch = .riscv32, .os_tag = .other, .cpu_features_add = features },
});

const apps = [_]String{
    "cat.c",
    "echo.c",
    "grep.c",
    "kill.c",
    "mkdir.c",
    "rm.c",
    "wc.c",
    "zombie.c",
    "primes.c",
    "find.c",
    "xargs.c",
    "trace.c",

    "init.zig",
    "sh.zig",
    "ls.zig",
    "ln.zig",
    "sleep.zig",
    "pingpong.zig",
    "shutdown.zig",

    "tests/forktest.c",
    "tests/stressfs.c",
    "tests/usertests.c",
    "tests/grind.c",
    "tests/bigfile.c",
    "tests/symlinktest.c",
    "tests/nettest.c",
};

const tools_platforms: []const Target.Query = &.{
    .{ .cpu_arch = .x86, .os_tag = .linux },
    .{ .cpu_arch = .x86_64, .os_tag = .linux },
    .{ .cpu_arch = .arm, .os_tag = .linux, .abi = .musleabi },
    .{ .cpu_arch = .aarch64, .os_tag = .linux },
    .{ .cpu_arch = .riscv32, .os_tag = .linux },
    .{ .cpu_arch = .riscv64, .os_tag = .linux },
    .{ .cpu_arch = .s390x, .os_tag = .linux },
    .{ .cpu_arch = .powerpc64le, .os_tag = .linux },
    .{ .cpu_arch = .loongarch64, .os_tag = .linux },

    .{ .cpu_arch = .x86, .os_tag = .windows },
    .{ .cpu_arch = .x86_64, .os_tag = .windows },
    .{ .cpu_arch = .aarch64, .os_tag = .windows },

    .{ .cpu_arch = .x86_64, .os_tag = .macos },
    .{ .cpu_arch = .aarch64, .os_tag = .macos },

    .{ .cpu_arch = .x86_64, .os_tag = .freebsd },
    .{ .cpu_arch = .aarch64, .os_tag = .freebsd },
    .{ .cpu_arch = .riscv64, .os_tag = .freebsd },
    .{ .cpu_arch = .powerpc64, .os_tag = .freebsd },
    .{ .cpu_arch = .powerpc64le, .os_tag = .freebsd },

    .{ .cpu_arch = .wasm32, .os_tag = .wasi },
};

pub fn build(b: *Build) !void {
    const now = try std.Io.Clock.now(.real, b.graph.io);

    const random = b.option(bool, "random", "Randomize CPU count, SBI, port and arch target") orelse false;
    var prng = std.Random.DefaultPrng.init(@bitCast(now.toSeconds()));

    const os_arch: OsArch = b.option(OsArch, "arch", "Target architecture for the OS kernel") orelse
        if (random) prng.random().enumValue(OsArch) else .riscv64;

    const os_target = b.resolveTargetQuery(os_arch_specs.get(os_arch));
    const optimize = b.standardOptimizeOption(.{});
    const strip = b.option(bool, "strip", "Removes symbols and sections from file");

    const use_sbi = b.option(bool, "sbi", "Enable OpenSBI firmware as BIOS") orelse
        if (random) prng.random().boolean() else false;

    const is_test = b.option(bool, "test", "Build kernel with unit tests") orelse false;
    const port_base = b.option(usize, "port", "Port base of nettest and gdb server") orelse
        if (random) prng.random().uintLessThanBiased(usize, 5000) else 0;

    // nettest ports
    const server_port = (port_base % 5000) + 25099;
    const fwd_port1 = (port_base % 5000) + 26579;
    const fwd_port2 = (port_base % 5000) + 31579;

    // gdb server port
    const gdb_port = (port_base % 5000) + 26002;

    // build kernel
    const os_module = b.createModule(.{
        .root_source_file = b.path("kernel/os.zig"),
    });

    const kernel_name = if (is_test) "test-kernel" else "kernel";
    const kernel = addKernelStep(b, .{
        .name = kernel_name,
        .target = os_target,
        .optimize = optimize,
        .strip = strip,
        .m_mode = !use_sbi,
        .is_test = is_test,
    });

    const kernel_path = kernel.getEmittedBin();
    const install_kernel = b.addInstallFile(kernel_path, kernel_name);
    b.getInstallStep().dependOn(&install_kernel.step);

    const kernel_tls = b.step("kernel", "Build and install kernel");
    kernel_tls.dependOn(&install_kernel.step);

    // build libc
    const libc = b.addLibrary(.{
        .name = "c",
        .root_module = b.createModule(.{
            .root_source_file = b.path("user/lib/c.zig"),
            .target = os_target,
            .optimize = optimize,
            .strip = strip,
            .code_model = .medium,
            .omit_frame_pointer = false,
            .sanitize_c = .off,
        }),
    });
    libc.link_function_sections = true;
    libc.link_data_sections = true;
    libc.link_gc_sections = true;

    libc.root_module.addImport("os", os_module);
    libc.addCSourceFiles(.{
        .root = b.path("user/lib/c"),
        .files = &.{ "ulib.c", "umalloc.c" },
    });

    libc.addIncludePath(b.path("user/include"));
    libc.installHeadersDirectory(b.path("user/include"), "", .{});

    const install_libc = b.addInstallArtifact(libc, .{});
    const libc_tls = b.step("libc", "Build and install libc");
    libc_tls.dependOn(&install_libc.step);

    // build user applications
    var rootfs = b.addWriteFiles();
    rootfs.step.name = "create rootfs";

    const is_rv32: u8 = if (os_target.result.cpu.arch == .riscv32) 1 else 0;
    for (apps) |app| {
        const exe = addAppStep(b, app, .{
            .target = os_target,
            .optimize = optimize,
            .strip = strip,
            .kernel = os_module,
            .libc = libc,
            .flags = &.{
                b.fmt("-DNET_TESTS_PORT={}", .{server_port}),
                b.fmt("-DIS_32BIT={}", .{is_rv32}),
            },
        });

        const subpath = if (std.mem.startsWith(u8, app, "tests/")) "tests" else "bin";
        _ = rootfs.addCopyFile(exe.getEmittedBin(), b.pathJoin(&.{ subpath, std.fs.path.stem(app) }));
    }

    // build mkfs
    const fs_module = b.createModule(.{
        .root_source_file = b.path("kernel/fs.zig"),
    });

    const mkfs = addMkfsStep(b, .{
        .fs_arch = os_target.result.cpu.arch,
        .target = b.graph.host,
        .optimize = optimize,
        .strip = strip,
        .fs = fs_module,
    });

    const tools_tls = b.step("tools", "Build and install tools");
    tools_tls.dependOn(&mkfs.step);

    // run mkfs to build fs.img
    const fs_size = b.option([]const u8, "fs-size", "Size of file system in blocks");
    const run_mkfs = b.addRunArtifact(mkfs.artifact);
    const fs_img = run_mkfs.addOutputFileArg("fs.img");
    _ = rootfs.addCopyFile(b.path("README.md"), "README");

    const script_opt = b.option(bool, "script", "Embed single-run test script in fs.img") orelse false;
    if (script_opt) {
        const script: []const u8 = if (b.args) |args|
            try std.mem.join(b.allocator, " ", args)
        else
            "forktest && stressfs && symlinktest && usertests";

        const test_and_exit = b.fmt("{s} && shutdown || shutdown 1", .{script});
        _ = rootfs.add(".shrc", test_and_exit);
    }

    run_mkfs.addDirectoryArg(rootfs.getDirectory());
    if (fs_size) |size| run_mkfs.addArg(size);
    _ = run_mkfs.captureStdErr(.{});

    const install_fs_img = b.addInstallFile(fs_img, "fs.img");
    b.getInstallStep().dependOn(&install_fs_img.step);

    const fs_tls = b.step("fs", "Build and install fs.img");
    fs_tls.dependOn(&install_fs_img.step);

    // runtime options
    const cpus = cpu: {
        const description = b.fmt("Number of cpus (1-{})", .{os.NCPU});
        const cpu_count = @min(os.NCPU, std.Thread.getCpuCount() catch 3);

        const option: u4 = b.option(u4, "cpus", description) orelse if (random) ra: {
            const int = prng.random().uintLessThanBiased(usize, cpu_count) + 1;
            break :ra @intCast(int);
        } else cpu_count;
        if (option > os.NCPU or option == 0) break :cpu cpu_count;
        break :cpu option;
    };

    const gdb_server = b.option(bool, "gdb", "Run QEMU with gdb server") orelse false;
    const gdb_args: []const String = if (gdb_server) &.{ "-gdb", b.fmt("tcp::{d}", .{gdb_port}), "-S" } else &.{};

    const custom_fs_path = b.option([]const u8, "fs-path", "Path to the custom fs.img");
    const fs_img_path = if (custom_fs_path) |path| b.path(path) else fs_img;

    // run xv6 in qemu
    const run_tls = b.step("run", "Run xv6 in QEMU");
    const run = addQemuStep(b, .{
        .arch = os_target.result.cpu.arch,
        .kernel = kernel_path,
        .fs_path = fs_img_path,
        .m_mode = !use_sbi,
        .cpus = cpus,
        .ports = .{ fwd_port1, fwd_port2 },
        .args = gdb_args,
    });

    run_tls.dependOn(&run.step);

    // debug with gdb
    const cross = std.process.getEnvVarOwned(b.allocator, "CROSS_COMPILE") catch "";
    const gdb_exe = std.process.getEnvVarOwned(b.allocator, "GDB") catch b.fmt("{s}gdb", .{cross});

    const gdb_tls = b.step("gdb", "Debug with gdb");
    const gdb = b.addSystemCommand(&.{ gdb_exe, "-q" });
    gdb.addArgs(&.{ "-ex", b.fmt("target remote : {}", .{gdb_port}) });
    gdb.addArgs(&.{ "-x", ".gdbinit" });
    gdb.addFileArg(kernel_path);

    gdb_tls.dependOn(&gdb.step);

    // display code information
    const objdump_tls = b.step("asm", "Display assembly code information");
    const objdump = b.addSystemCommand(&.{ b.fmt("{s}objdump", .{cross}), "-S" });
    objdump.addFileArg(kernel_path);

    objdump_tls.dependOn(&objdump.step);

    // translates addresses of stack trace
    const addr2line_tls = b.step("addr2line", "Translates addresses of stack trace");
    const addr2line = b.addSystemCommand(&.{b.fmt("{s}addr2line", .{cross})});
    addr2line.addArgs(&.{"--exe"});
    addr2line.addFileArg(kernel_path);

    if (b.args) |args| {
        addr2line.addArgs(args);
    }

    addr2line_tls.dependOn(&addr2line.step);

    // build and run test tool
    const net_option = b.addOptions();
    net_option.addOption(usize, "server", server_port);
    net_option.addOption(usize, "fwd1", fwd_port1);
    net_option.addOption(usize, "fwd2", fwd_port2);

    const nettest = addNettestStep(b, .{
        .target = b.graph.host,
        .optimize = optimize,
        .strip = strip,
        .option = net_option,
    });

    tools_tls.dependOn(&nettest.step);

    const run_nettest = b.addRunArtifact(nettest.artifact);
    if (b.args) |args| {
        run_nettest.addArgs(args);
    }

    const nettest_tls = b.step("nettest", "Run nettest tool");
    nettest_tls.dependOn(&run_nettest.step);

    // release builds of tools for multiple targets
    const release_tools = b.step("release/tools", "Build tools for multiple platforms");
    for (tools_platforms) |t| {
        const suffix = b.fmt("{t}-{t}", .{ t.os_tag.?, t.cpu_arch.? });
        const target = b.resolveTargetQuery(t);

        const mk = addMkfsStep(b, .{
            .name = b.fmt("mkfs-{s}", .{suffix}),
            .target = target,
            .fs = fs_module,
            .fs_arch = os_target.result.cpu.arch,
        });
        release_tools.dependOn(&mk.step);

        const nt = addNettestStep(b, .{
            .name = b.fmt("nettest-{s}", .{suffix}),
            .target = target,
            .option = net_option,
        });
        release_tools.dependOn(&nt.step);
    }
}

const KernelOptions = struct {
    name: String,
    target: Build.ResolvedTarget,
    optimize: std.builtin.OptimizeMode,
    strip: ?bool = null,
    is_test: bool = false,
    m_mode: bool = true,
};

fn addKernelStep(b: *Build, opt: KernelOptions) *Step.Compile {
    // riscv registers module
    const riscv_dep = b.dependency("riscv", .{
        .target = opt.target,
        .optimize = opt.optimize,
    });

    // libfdt module
    const dtb_dep = b.dependency("dtb", .{
        .target = opt.target,
        .optimize = opt.optimize,
    });

    // options
    const options = b.addOptions();
    options.addOption(bool, "m_mode", opt.m_mode);

    // kernel root module
    const kernel_module = b.createModule(.{
        .root_source_file = b.path("kernel/start.zig"),
        .target = opt.target,
        .optimize = opt.optimize,
        .strip = opt.strip,
        .code_model = .medium,
        .omit_frame_pointer = false,
    });

    kernel_module.addOptions("options", options);
    kernel_module.addImport("dtb", dtb_dep.module("dtb"));
    kernel_module.addImport("riscv", riscv_dep.module("riscv"));

    // build kernel binary
    const kernel = if (opt.is_test)
        b.addTest(.{
            .name = opt.name,
            .root_module = kernel_module,
            .test_runner = .{
                .path = b.path("tools/test_runner.zig"),
                .mode = .simple,
            },
        })
    else
        b.addExecutable(.{
            .name = opt.name,
            .root_module = kernel_module,
        });

    if (opt.is_test) {
        kernel.root_module.addImport("kernel", kernel_module);
    }

    const linker = b.path("kernel/arch/kernel.ld");
    kernel.setLinkerScript(linker);
    kernel.entry = .{ .symbol_name = "_entry" };
    kernel.image_base = if (opt.m_mode) 0x8000_0000 else 0x8020_0000;
    kernel.compress_debug_sections = .zlib;

    return kernel;
}

const QemuOptions = struct {
    arch: Target.Cpu.Arch,
    kernel: Build.LazyPath,
    fs_path: Build.LazyPath,
    m_mode: bool = true,
    cpus: usize,
    ports: struct { usize, usize },
    args: []const String = &.{},
};

fn addQemuStep(b: *Build, opt: QemuOptions) *Step.Run {
    const qemu_cmd = b.fmt("qemu-system-{t}", .{opt.arch});
    const arch_bit: u16 = if (opt.arch == .riscv64) 64 else 32;
    const packets_path = b.cache_root.join(b.allocator, &.{"packets.pcap"}) catch unreachable;

    const qemu_args = [_][]const u8{
        "-machine",   "virt",
        "-cpu",       b.fmt("rv{d},zihintpause=true", .{arch_bit}),
        "-bios",      if (opt.m_mode) "none" else "default",
        "-m",         "128M",
        "-smp",       b.fmt("{}", .{opt.cpus}),

        // https://www.qemu.org/docs/master/system/invocation.html#hxtool-1
        "-global",    "virtio-mmio.force-legacy=false",
        "-device",    "virtio-blk-device,drive=x0,bus=virtio-mmio-bus.0",

        // https://www.qemu.org/docs/master/system/devices/net.html#using-the-user-mode-network-stack
        "-netdev",    b.fmt("user,id=net0,hostfwd=udp::{}-:2000,hostfwd=udp::{}-:2001", opt.ports),
        "-object",    b.fmt("filter-dump,id=net0,netdev=net0,file={s}", .{packets_path}),
        "-device",    "e1000,netdev=net0,bus=pcie.0",

        "-nographic",
    };

    const run = b.addSystemCommand(&.{qemu_cmd});
    run.addArgs(&qemu_args);
    if (opt.args.len > 0) run.addArgs(opt.args);

    run.addArgs(&.{"-kernel"});
    run.addFileArg(opt.kernel);

    run.addArgs(&.{"-drive"});
    run.addPrefixedFileArg("id=x0,if=none,format=raw,file=", opt.fs_path);

    return run;
}

const AppOptions = struct {
    target: Build.ResolvedTarget,
    optimize: std.builtin.OptimizeMode,
    strip: ?bool = null,
    kernel: *Build.Module,
    libc: *Step.Compile,
    flags: []const []const u8 = &.{},
};

fn addAppStep(b: *Build, file_name: String, opt: AppOptions) *Step.Compile {
    const file_ext = std.fs.path.extension(file_name);
    const file_type: enum { c, zig } = if (std.mem.eql(u8, file_ext, ".zig")) .zig else .c;
    const app_path = b.pathJoin(&.{ "user/", file_name });

    const app = b.addExecutable(.{
        .name = std.fs.path.stem(file_name),
        .root_module = b.createModule(.{
            .root_source_file = if (file_type == .c)
                null
            else
                b.path(app_path),
            .target = opt.target,
            .optimize = opt.optimize,
            .strip = opt.strip,
            .code_model = .medium,
            .omit_frame_pointer = false,
            .no_builtin = true,
        }),
    });

    if (file_type == .c) {
        app.root_module.sanitize_c = .off;
        app.linkLibrary(opt.libc);
        app.addCSourceFile(.{ .file = b.path(app_path), .flags = opt.flags });
    } else {
        app.root_module.addImport("os", opt.kernel);

        if (std.mem.eql(u8, file_name, "sh.zig")) {
            const shell_dep = b.dependency("shell", .{
                .target = opt.target,
                .optimize = opt.optimize,
            });
            app.root_module.addImport("Ast", shell_dep.module("shell"));
            app.linkLibrary(opt.libc);
        }
    }

    app.link_gc_sections = true;
    app.compress_debug_sections = .zlib;
    app.entry = .{ .symbol_name = "main" };
    app.setLinkerScript(b.path("user/app.ld"));

    return app;
}

const MkfsOptions = struct {
    name: String = "mkfs",
    fs_arch: Target.Cpu.Arch,
    target: Build.ResolvedTarget,
    optimize: std.builtin.OptimizeMode = .ReleaseSmall,
    strip: ?bool = null,
    fs: *Build.Module,
};

fn addMkfsStep(b: *Build, opt: MkfsOptions) *Step.InstallArtifact {
    const mkfs = b.addExecutable(.{
        .name = opt.name,
        .root_module = b.createModule(.{
            .root_source_file = b.path("tools/mkfs.zig"),
            .target = opt.target,
            .optimize = opt.optimize,
            .single_threaded = true,
        }),
    });

    const mkfs_option = b.addOptions();
    mkfs_option.addOption(String, "arch", @tagName(opt.fs_arch));
    mkfs.root_module.addOptions("target", mkfs_option);
    mkfs.root_module.addImport("fs", opt.fs);

    return b.addInstallArtifact(mkfs, .{
        .dest_dir = .{ .override = .{ .custom = "tools" } },
    });
}

const NettestOptions = struct {
    name: String = "nettest",
    target: Build.ResolvedTarget,
    optimize: std.builtin.OptimizeMode = .ReleaseSmall,
    strip: ?bool = null,
    option: *Step.Options,
};

fn addNettestStep(b: *Build, opt: NettestOptions) *Step.InstallArtifact {
    const nettest = b.addExecutable(.{
        .name = opt.name,
        .root_module = b.createModule(.{
            .root_source_file = b.path("tools/nettest.zig"),
            .target = opt.target,
            .optimize = opt.optimize,
        }),
    });

    if (opt.target.result.os.tag == .windows) {
        nettest.root_module.link_libc = true;
        nettest.linkSystemLibrary("ws2_32");
    }

    nettest.root_module.addOptions("port", opt.option);

    return b.addInstallArtifact(nettest, .{
        .dest_dir = .{ .override = .{ .custom = "tools" } },
    });
}
